import pandas as pd


"""将txt每行作为一个单元格转换为只有一列的表格"""
raw = pd.read_csv(
    "../novel.txt", # 读取文件
    names = ['txt'], # 设置列名
    sep ='aaa',  # 分隔符，必须设置为文章中没有出现过的字符串
    encoding ="GBK", engine='python'
)
# print(len(raw)) # 行数
# print(raw) # 小说数据框

"""对raw数据框的txt列的每一行调用函数新建其他列"""
raw['char1st'] = raw.txt.apply(lambda line: line[0])
raw['index'] = raw.txt.apply(lambda line: line.find('回 '))
raw['length'] = raw.txt.apply(len)
raw['ch_num'] = 0 # 将char_num设置为整数，防止默认设置为小数

"""章节判断"""
ch_num = 0
# 遍历小说所有行，并注明这个单元格所属的章节号
for i in range(len(raw)):
    if raw['char1st'][i] == "第" and raw['index'][i] > 0 and raw['length'][i] < 30 :
        # 查找有 "第**回 " 并且文本数少于30的行，一般每章标题不会超过30字
        ch_num += 1
    if ch_num >= 40 and raw['txt'][i] == "附录一：成吉思汗家族" :
        # 正文总共40回，这里只需要正文内容，不需要前言附录等
        char_num = 0
    raw.loc[i, 'ch_num'] = ch_num

"""删除辅助列"""
del raw['char1st']
del raw['index']
del raw['length']

# print(raw.head(50))

# 按照ch_num列分组
rawgrp = raw.groupby('ch_num')
# 拼接同一组的内容为字符串
chapter = rawgrp.agg('sum')
# 过滤掉ch_num为0的章节
chapter = chapter[chapter.index != 0]
# print(chapter.txt[2])
# print(chapter)

def chapter_n(n: int) -> str:
    return str(chapter.txt[n])

def download_ch(n):
    with open(f'../notebook/ch-{n}.txt', 'w', encoding='utf-8') as f:
        f.write(chapter_n(n))

if __name__ == '__main__':
    # download_ch(2)
    pass